#!/usr/bin/env python3

import streamlit as st
import random
import os
import subprocess
from .tool_register import ToolRegistry
from .net.get_topo import get_topo, get_weighted_topo
from .net.get_switch_stats import get_switch_stats
from .net.get_path_delay import get_path_delay
from .utils import st_get_api
import streamlit.components.v1 as components
import matplotlib.pyplot as plt
import networkx as nx
import matplotlib.image as mpimg
from pathlib import Path
import re
import requests

tool_registry = ToolRegistry()

@tool_registry.tool
def search(query: str) -> str:
    """Look up things online."""
    return "Some results"

@tool_registry.tool
def plus(a: float, b: float) -> float:
    """Perform a basic addition of two numbers."""
    return a + b

@tool_registry.tool
def minus(a: float, b: float) -> float:
    """Perform a basic minus of two numbers."""
    return a - b

@tool_registry.tool
def times(a: float, b: float) -> float:
    """Perform a basic multiply of two numbers."""
    return a * b

@tool_registry.tool
def divide(a: float, b: float) -> float:
    """Perform a basic divide of two numbers."""
    return a / b

@tool_registry.tool
def ping(host: str) -> str:
    """Perform a basic ping to 'host', return response times in milliseconds."""
    try:
        # Run the ping command and capture the output
        command = ["ping", "-c", "4", host]
        result = subprocess.run(command, capture_output=True, text=True, check=True)

        # # Extract the relevant output (response times)
        # output_lines = result.stdout.splitlines()
        # response_times = []

        # for line in output_lines:
        #     if "time=" in line:
        #         # Extract the time value from the line
        #         time_part = line.split("time=")[-1]
        #         response_time = time_part.split(" ")[0]  # Get the value before the 'ms'
        #         response_times.append(response_time)

        # # Format the response times for output
        # formatted_response = ", ".join(response_times)
        # return f"Response times (ms): {formatted_response}"
        return result

    except subprocess.CalledProcessError as e:
        return f"Error pinging {host}: {e.stderr.strip()}"

@tool_registry.tool
def further_thought():
    """Perform further thought of assistant when current Response is not the final answer"""
    return

@tool_registry.tool
def sdn_init_quality():
    """This is the API of SDN control plane you are operating on,
    this function will send a request to the control plane to start a controller with delay detector.
    the routing path will be determined by the shortest delay"""
    end_point = '/load/quality'
    url, time_out = st_get_api(end_point)
    response = requests.post(url)

@tool_registry.tool
def sdn_init_hop():
    """This is the API of SDN control plane you are operating on,
    this function will send a request to the control plane to start a controller without delay detector.
    the routing path will be determined by the minimal hops"""
    end_point = '/load/hop'
    url, time_out = st_get_api(end_point)
    response = requests.post(url)

@tool_registry.tool
def sdn_get_switch_stats(switch_idx: int) -> str:
    """This is the API of SDN control plane you are operating on,
    this function will return the switch status information with queried index (switch_idx)."""
    stats, df = get_switch_stats(switch_idx)
    st.session_state.messages.append({"role": "table", "content": df})
    st.table(df)
    return stats

@tool_registry.tool
def sdn_get_switch_path_delay() -> str:
    """This is the API of SDN control plane you are operating on,
    this function will return the detected path delay between switches,
    which is echo_delay - LLDP_delay. The key of dict is structure like
    'src-src_port-dst'."""
    data, df = get_path_delay()
    st.session_state.messages.append({"role": "table", "content": df})
    st.table(df)
    return data

@tool_registry.tool
def sdn_get_switch_topo() -> str:
    """This is the API of SDN control plane you are operating on,
    this function will only return the topology information about switches."""
    raw, dot = get_topo()
    st.graphviz_chart(dot.source, use_container_width=False)
    st.caption('SDN Switch Topology')
    st.session_state.messages.append({"role": "graphviz", "content": dot})
    return raw

@tool_registry.tool
def sdn_get_gui_topo() -> str:
    """This is the API of SDN control plane you are operating on,
    this function will solely display the sdn switch topology to the user and return a
    'done' message ."""

    end_point = ''
    url, time_out = st_get_api(end_point)
    components.iframe(url, height=400, width=650, scrolling=True)
    message = 'done'
    return message

@tool_registry.tool
def sdn_get_weighted_topo() -> str:
    """This is the API of SDN control plane you are operating on,
    this function will return the information about all topology,
    devices with only mac address are the hosts, the edge weights
    are the delay."""
    raw, G = get_weighted_topo()
    fig, ax = plt.subplots(figsize=(10, 6))
    # pos = nx.spectral_layout(G)
    pos = nx.spring_layout(G, k=0.15)  # positions for all nodes

    # Load the custom node images
    path = Path(__file__).parent
    switch_img = mpimg.imread(path / "net/img/switch.png")
    host_img = mpimg.imread(path / "net/img/host.png")

    # Create a list for node colors and edge widths
    node_images = []
    edge_widths = [d['weight']*10000 for _, _, d in G.edges(data=True)] # multiply by 10000 to make line width visible

    # Associate image with each node
    for node in G.nodes(data=True):
        if isinstance(node[0], int):
            G.nodes[node[0]]["image"] = switch_img
        else:
            if re.match('^([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}$', node[0]):
                G.nodes[node[0]]["image"] = host_img

    # Draw the network
    # nx.draw_networkx_edges(G, pos, width=edge_widths)
    nx.draw_networkx_edges(G, pos, edge_color='royalblue')
    # nx.draw_networkx_labels(G, pos)
    offset = -0.15
    for node, (x, y) in pos.items():
        plt.text(x, y + offset, str(node), ha='center')
    # Transformations for the placement of images
    tr_figure = ax.transData.transform
    tr_axes = fig.transFigure.inverted().transform

    icon_size = (ax.get_xlim()[1] - ax.get_xlim()[0]) * 0.025
    icon_center = icon_size / 2.0

    # Add the respective image to each node
    for n in G.nodes:
        xf, yf = tr_figure(pos[n])
        xa, ya = tr_axes((xf, yf))
        a = plt.axes([xa - icon_center, ya - icon_center, icon_size, icon_size])
        a.imshow(G.nodes[n]["image"])
        a.axis("off")

    # Set title and remove axis
    ax.set_title("Network Topology with Delay")
    ax.axis('off')  # Turn off the axis

    # Use Streamlit to display the figure
    st.pyplot(fig)
    st.session_state.messages.append({"role": "fig", "content": fig})

    return raw

# @tool_registry.tool
# def sdn_get_weighted_topo() -> str:
#     """This is the API of SDN control plane you are operating on,
#     this function will return the information about all topology,
#     devices with only mac address are the hosts, the edge weights
#     are the delay."""
#     raw, G = get_weighted_topo()
#     fig, ax = plt.subplots(figsize=(10, 6))

#     pos = nx.spring_layout(G)  # positions for all nodes
#     # nx.draw_networkx_nodes(G, pos, ax=ax)
#     # nx.draw_networkx_edges(G, pos, ax=ax, arrows=True)

#     # # Draw labels
#     # labels = {node['id']: node['id'] for node in network_data['nodes']}
#     # nx.draw_networkx_labels(G, pos, labels, ax=ax)

#     # # Display edge weights
#     # edge_labels = nx.get_edge_attributes(G, 'weight')
#     # nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax)
#     nx.draw(G, with_labels=True)
#     # Set title and remove axis
#     ax.set_title("Network Topology with Delay")
#     ax.axis('off')  # Turn off the axis

#     # Use Streamlit to display the figure
#     st.pyplot(fig)

#     st.session_state.messages.append({"role": "fig", "content": fig})

#     return raw

if __name__ == "__main__":
    # Example: show available tools and their fns
    print(tool_registry.tools)

    # Example: get arguments
    for tool in tool_registry.tools.values():
        print(tool.args)

    # Example: get text description of all tools
    print(tool_registry.render_text_description())

    # Example: Use a registered tool
    result = tool_registry.tools["search"]("example query")
    print(f"Search result: {result}")

    # Example: Add a tool
    @tool_registry.tool
    def new_tool(query: str) -> str:
        """One new tool."""
        return "Some results"
    print(tool_registry.render_text_description())

